import serial
import serial.tools.list_ports
import threading
import time
import json
from flask import Flask, jsonify, request
import crcmod.predefined  # 确保安装：pip install crcmod
from flask_cors import CORS
from datetime import datetime

app = Flask(__name__)
CORS(app)

# 全局状态管理
class GlobalState:
    def __init__(self):
        self.lock = threading.Lock()
        self.system_status = "初始化中..."
        self.current_angle = None  # 云台当前角度
        self.running = False  # 是否正在测量
        self.current_round = 0  # 当前轮次
        self.config = {
            "stepper_cycles": 1000,  # 步进电机运行周期
            "move_distance_mm": 10,  # 每轮移动距离
            "total_rounds": 50,  # 总轮次
            "baudrate": 115200,
            "timeout": 0.5
        }
        self.stepper = None  # 步进电机实例
        self.gimbal = None  # 云台电机实例

state = GlobalState()

# -------------------------- 基础工具类 --------------------------
class CRC16:
    """CRC16校验工具类"""
    @staticmethod
    def calculate(data):
        crc16 = crcmod.predefined.Crc('modbus')
        crc16.update(data)
        return crc16.digest()

# -------------------------- 步进电机控制类 --------------------------
class StepperMotor:
    def __init__(self, motor_id=3, port=None, baudrate=115200, timeout=0.5):
        self.motor_id = motor_id
        self.baudrate = baudrate
        self.timeout = timeout
        self.ser = None
        self.lock = threading.Lock()
        self.connected = False
        self._connect(port)

    def _connect(self, port):
        """连接步进电机串口"""
        if not port:
            ports = list(serial.tools.list_ports.comports())
            for p in ports:
                try:
                    self.ser = serial.Serial(p.device, self.baudrate, timeout=self.timeout)
                    time.sleep(0.2)
                    self.connected = True
                    print(f"步进电机连接成功: {p.device}")
                    return
                except Exception as e:
                    continue
            print("步进电机未找到可用串口")
            return
        try:
            self.ser = serial.Serial(port, self.baudrate, timeout=self.timeout)
            time.sleep(0.2)
            self.connected = True
            print(f"步进电机连接成功: {port}")
        except Exception as e:
            print(f"步进电机连接失败: {e}")

    def _send_modbus_command(self, register, value):
        """发送Modbus命令（带重试机制）"""
        if not self.connected or not self.ser:
            print("[Stepper] 未连接，无法发送命令")
            return False

        with self.lock:
            retry = 3
            cmd = bytearray()
            cmd.append(self.motor_id)
            cmd.append(0x06)  # 写单个寄存器
            cmd.extend(register.to_bytes(2, byteorder='big'))
            cmd.extend(value.to_bytes(2, byteorder='big'))
            crc = CRC16.calculate(cmd)
            cmd.extend(crc)

            while retry > 0:
                try:
                    self.ser.reset_input_buffer()
                    self.ser.reset_output_buffer()
                    self.ser.write(cmd)
                    self.ser.flush()
                    time.sleep(0.08)  # 延长等待响应时间
                    resp = self.ser.read(8)  # Modbus写响应固定8字节

                    if len(resp) == 8:
                        # 验证响应是否与命令一致
                        if (resp[0] == self.motor_id and resp[1] == 0x06 and
                            int.from_bytes(resp[2:4], 'big') == register and
                            int.from_bytes(resp[4:6], 'big') == value):
                            return True
                        else:
                            print(f"[Stepper] 响应内容不匹配: {resp.hex()}")
                    else:
                        print(f"[Stepper] 写寄存器未返回预期长度: reg=0x{register:04X} val=0x{value:04X} resp_len={len(resp)}")
                    retry -= 1
                    time.sleep(0.1)  # 重试间隔
                except Exception as e:
                    print(f"[Stepper] 发送命令异常(重试{3-retry}次): {e}")
                    retry -= 1
                    time.sleep(0.1)

            print(f"[Stepper] 命令最终失败: reg=0x{register:04X} val=0x{value:04X}")
            return False

    def rotate_cycles(self, cycles):
        """设置旋转周期"""
        return self._send_modbus_command(0x6001, cycles)

    def rotate_start(self, cycles):
        """启动旋转（先设置周期再触发运行）"""
        if not self.rotate_cycles(cycles):
            print("[Stepper] 设置周期失败")
            return False
        # 触发PRO运行（寄存器0x6002，值0x0010）
        return self._send_modbus_command(0x6002, 0x0010)

    def stop(self):
        """停止步进电机"""
        return self._send_modbus_command(0x6002, 0x0000)

    def close(self):
        """关闭串口"""
        if self.ser and self.ser.is_open:
            self.ser.close()
        self.connected = False

# -------------------------- 云台电机控制类 --------------------------
class HaitaiMotor:
    """海泰云台电机控制类（统一角度计算逻辑）"""
    def __init__(self, port=None, baudrate=115200, timeout=0.5):
        self.baudrate = baudrate
        self.timeout = timeout
        self.ser = None
        self.lock = threading.Lock()
        self.connected = False
        self.counts_per_circle = 131072  # 电机每圈计数（根据实际型号调整）
        self._connect(port)

    def _connect(self, port):
        """连接云台电机串口"""
        if not port:
            ports = list(serial.tools.list_ports.comports())
            for p in ports:
                try:
                    self.ser = serial.Serial(p.device, self.baudrate, timeout=self.timeout)
                    time.sleep(0.2)
                    self.connected = True
                    print(f"云台电机连接成功: {p.device}")
                    return
                except Exception as e:
                    continue
            print("云台电机未找到可用串口")
            return
        try:
            self.ser = serial.Serial(port, self.baudrate, timeout=self.timeout)
            time.sleep(0.2)
            self.connected = True
            print(f"云台电机连接成功: {port}")
        except Exception as e:
            print(f"云台电机连接失败: {e}")

    def _send_command(self, cmd_code, data, read_len=0, post_delay=0.1):
        """发送云台电机命令"""
        if not self.connected or not self.ser:
            print("[Gimbal] 未连接，无法发送命令")
            return None

        with self.lock:
            try:
                # 命令格式：帧头(0xAA) + 地址(0x01) + 长度 + 命令码 + 数据 + 校验和
                frame = bytearray([0xAA, 0x01])
                frame.append(len(data) + 1)  # 数据长度+命令码长度
                frame.append(cmd_code)
                frame.extend(data)
                # 计算校验和（异或）
                checksum = 0
                for b in frame[1:]:
                    checksum ^= b
                frame.append(checksum)

                self.ser.reset_input_buffer()
                self.ser.reset_output_buffer()
                self.ser.write(frame)
                self.ser.flush()
                time.sleep(post_delay)

                if read_len > 0:
                    resp = self.ser.read(read_len)
                    # 验证响应帧头
                    if resp and resp[0] == 0xAA:
                        return resp
                    else:
                        print(f"[Gimbal] 响应帧头错误: {resp.hex() if resp else 'None'}")
                return None
            except Exception as e:
                print(f"[Gimbal] 发送命令异常: {e}")
                return None

    def set_origin(self):
        """设置当前位置为原点"""
        resp = self._send_command(0x10, b'\x00\x00\x00\x00', read_len=6)
        if resp and len(resp) == 6 and resp[3] == 0x00:
            print("原点设置成功")
            return True
        else:
            print(f"原点设置失败: {resp.hex() if resp else 'None'}")
            return False

    def get_current_angle(self):
        """读取当前角度（统一计算逻辑）"""
        try:
            resp = self._send_command(0x2F, b'', read_len=15, post_delay=0.1)
        except Exception as e:
            print(f"[Gimbal] 读取角度出错: {e}")
            return None

        if resp and len(resp) == 15:
            try:
                # 解析单圈角度（5-6字节）
                single_turn_raw = (resp[5] | (resp[6] << 8)) & 0x3FFF  # 14位有效数据
                single_turn_angle = single_turn_raw * (360.0 / self.counts_per_circle)
                
                # 解析多圈角度（7-10字节）
                multi_turn_raw = (resp[7] | (resp[8] << 8) | (resp[9] << 16) | (resp[10] << 24))
                # 处理负数（补码转换）
                if multi_turn_raw & 0x80000000:
                    multi_turn_raw -= 0x100000000
                multi_turn_angle = multi_turn_raw * (360.0 / self.counts_per_circle)
                
                # 统一角度计算：多圈角度 + 单圈角度（取模360避免重复）
                total_angle = multi_turn_angle + (single_turn_angle % 360)
                return round(total_angle, 2)  # 保留2位小数
            except Exception as e:
                print(f"[Gimbal] 解析角度失败: {e} | 应答数据: {resp.hex()}")
                return None
        else:
            print(f"[Gimbal] 应答长度错误（需15字节）: 实际={len(resp) if resp else 0}")
            return None

    def close(self):
        """关闭串口"""
        if self.ser and self.ser.is_open:
            self.ser.close()
        self.connected = False

# -------------------------- 测量任务 --------------------------
def angle_monitor_task():
    """云台角度监控线程（实时更新角度）"""
    while True:
        if state.gimbal and state.gimbal.connected:
            angle = state.gimbal.get_current_angle()
            with state.lock:
                state.current_angle = angle
        time.sleep(0.1)  # 100ms刷新一次

def measurement_task():
    """主测量任务"""
    with state.lock:
        state.running = True
        state.system_status = "开始回归零位..."
    
    # 回归零位并设置原点
    if state.gimbal and state.gimbal.connected:
        # 此处根据实际零位回归逻辑修改（示例：假设已手动回归零位）
        time.sleep(2)  # 模拟回归零位时间
        with state.lock:
            state.system_status = "设置当前位置为原点..."
        if state.gimbal.set_origin():
            with state.lock:
                state.system_status = "原点设置成功"
        else:
            with state.lock:
                state.system_status = "原点设置失败"
            state.running = False
            return
    else:
        with state.lock:
            state.system_status = "云台电机未连接"
        state.running = False
        return

    time.sleep(1)
    with state.lock:
        state.system_status = "已到达零位"

    # 开始多轮测量
    for i in range(state.config["total_rounds"]):
        with state.lock:
            state.current_round = i + 1
            state.system_status = f"第{state.current_round}轮: 移动{state.config['move_distance_mm']}mm"
        
        # 控制步进电机移动
        if state.stepper and state.stepper.connected:
            # 尝试启动步进电机（失败则重试初始化）
            if not state.stepper.rotate_start(state.config["stepper_cycles"]):
                with state.lock:
                    state.system_status = f"第{state.current_round}轮：步进电机启动失败，尝试重新初始化..."
                # 重新初始化步进电机
                if not state.stepper.rotate_cycles(state.config["stepper_cycles"]):
                    with state.lock:
                        state.system_status = f"第{state.current_round}轮：步进电机初始化失败"
                    break
                # 再次尝试启动
                if not state.stepper.rotate_start(state.config["stepper_cycles"]):
                    with state.lock:
                        state.system_status = f"第{state.current_round}轮：步进电机启动失败"
                    break
        
        # 等待移动完成（根据实际移动时间调整）
        move_time = state.config["stepper_cycles"] * 0.01  # 示例计算
        time.sleep(max(move_time, 1))  # 最小等待1秒

        # 停止步进电机
        if state.stepper and state.stepper.connected:
            state.stepper.stop()

        # 每轮测量间隔
        time.sleep(0.5)

    # 测量完成
    with state.lock:
        state.running = False
        state.system_status = f"测量完成（共{state.current_round}轮）"

# -------------------------- Flask API --------------------------
@app.route("/api/status", methods=["GET"])
def get_status():
    """获取系统状态和云台角度"""
    with state.lock:
        return jsonify({
            "running": state.running,
            "current_round": state.current_round,
            "current_angle": state.current_angle,
            "system_status": state.system_status,
            "stepper_connected": state.stepper.connected if state.stepper else False,
            "gimbal_connected": state.gimbal.connected if state.gimbal else False
        })

@app.route("/api/start", methods=["POST"])
def start_measurement():
    """启动测量任务"""
    if state.running:
        return jsonify({"code": 1, "msg": "测量已在运行中"}), 200
    
    # 解析请求参数（可选）
    data = request.get_json() or {}
    if "total_rounds" in data:
        state.config["total_rounds"] = int(data["total_rounds"])
    if "move_distance_mm" in data:
        state.config["move_distance_mm"] = int(data["move_distance_mm"])
    
    # 启动测量线程
    threading.Thread(target=measurement_task, daemon=True).start()
    return jsonify({"code": 0, "msg": "测量已启动"}), 200

@app.route("/api/stop", methods=["POST"])
def stop_measurement():
    """停止测量任务"""
    with state.lock:
        state.running = False
        state.system_status = "测量已停止"
    
    # 停止步进电机
    if state.stepper and state.stepper.connected:
        state.stepper.stop()
    
    return jsonify({"code": 0, "msg": "测量已停止"}), 200

@app.route("/api/reset", methods=["POST"])
def reset_system():
    """重置系统状态"""
    with state.lock:
        state.current_round = 0
        state.current_angle = None
        state.system_status = "已重置"
        state.running = False
    
    # 停止电机
    if state.stepper and state.stepper.connected:
        state.stepper.stop()
    
    return jsonify({"code": 0, "msg": "系统已重置"}), 200

# -------------------------- 程序入口 --------------------------
if __name__ == "__main__":
    try:
        # 初始化电机
        state.stepper = StepperMotor(motor_id=3)
        state.gimbal = HaitaiMotor()

        # 启动角度监控线程
        threading.Thread(target=angle_monitor_task, daemon=True).start()

        # 启动Flask服务
        print("后端服务启动成功，访问 http://127.0.0.1:5000")
        app.run(host="0.0.0.0", port=5000, debug=False)  # 生产环境关闭debug
    except KeyboardInterrupt:
        print("程序退出，关闭资源...")
    finally:
        # 关闭串口连接
        if state.stepper:
            state.stepper.close()
        if state.gimbal:
            state.gimbal.close()